/*
 * Copyright (c) Extreme Vision Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef YOLOx_PREPOSTPROCESSOR_H
#define YOLOx_PREPOSTPROCESSOR_H

#include "opencv2/core.hpp"
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include "evdeploy/core/ev_common_data.h"
#include "evdeploy/cv/common/common.h"

namespace ev
{
namespace vision
{

class YOLOxPreprocessor
{
    public:        
        YOLOxPreprocessor();

        ~YOLOxPreprocessor();        
        /**
         * @brief preprocessing for YOLOx
         * @param[in] in_mat: input cv::Mat
         * @param[in] out_width: resized Mat Size(out_width, out_width)
         * @param[in] normalize: true for normalize to 1, false for without normalization, default false
         * @param[in] bgr2rgb: true for bgr2rgb, false for bgr format, default false
         * @param[out] out_mat: output cv::Mat
         * @return EVStatus
         */
        EVStatus Run(cv::Mat& in_mat, cv::Mat& out_mat, int out_width, bool normalize=false, bool bgr2rgb=false);
        /**
         * @brief get scale from preprocessing, used for postprocessing, e.g., compute the coordinates in original image
         * @return scale
         */
        float GetScale();
       
    private:             
        float m_scale;
};

struct GridAndStride
{
    int grid0;
    int grid1;
    int stride;
};

class YOLOxPostprocessor
{
    public:        
        YOLOxPostprocessor();

        ~YOLOxPostprocessor();
        /**
         * @brief postprocessing for YOLOv5
         * @param[in] out: output data from runtime inference
         * @param[in] scale: scale from preprocessing
         * @param[in] thresh: thresh for detection confidence
         * @param[in] model_resolution: model_resolution for inference, used to compute coordinates in original image 
         * @param[in] img_w: original image width
         * @param[in] img_h: original image height
         * @param[in] nms_thresh: thresh for nms
         * @param[in] class_aware whether to distinguish labels in nms
         * @param[out] objects: detection output wrapped in BoxInfo array 
         * @return EVStatus
         */
        EVStatus Run(EVMatData *out, std::vector<ev::vision::BoxInfo> &objects, float scale, float thresh, int model_resolution, const int img_w, const int img_h, float nms_thresh = 0.45, bool class_aware = false);

};


}//namespace vision
}//namespace ev

#endif 
